"""
自注意力机制 (Self-Attention Mechanism)
这是Transformer的核心组件，让AI能够理解句子中词与词之间的关系

就像人在理解句子时，会关注不同词汇之间的联系：
"小明在北京大学读书" - 当理解"读书"时，会联想到"小明"和"大学"
"""

import numpy as np
from typing import Optional, Tuple
from adam import Adam

# 常量
EMBEDDING_DIM = 32  # 词向量维度

class SelfAttention:
    """
    自注意力层 - AI的"重点关注器"
    
    核心思想：
    1. 对于句子中的每个词，计算它应该关注其他哪些词
    2. 根据关注程度，融合相关词汇的信息
    3. 让AI更好地理解词汇的上下文含义
    
    比如在句子"猫坐在垫子上"中：
    - 处理"坐"时，会重点关注"猫"（主语）和"垫子"（地点）
    - 处理"垫子"时，会关注"猫"和"坐"的关系
    """
    
    def __init__(self, embedding_dim: int):
        """
        初始化自注意力层
        embedding_dim: 词向量的维度（每个词用多少个数字表示）
        """
        self.embedding_dim = embedding_dim
        
        # Xavier/He初始化：一种让训练更稳定的权重初始化方法
        std = np.sqrt(2.0 / embedding_dim)
        
        # 创建三个权重矩阵：Q(查询)、K(键)、V(值)
        # 就像三个不同的"镜头"，从不同角度观察同一个词
        self.w_q = np.random.normal(0.0, std, (embedding_dim, embedding_dim))  # Query权重
        self.w_k = np.random.normal(0.0, std, (embedding_dim, embedding_dim))  # Key权重  
        self.w_v = np.random.normal(0.0, std, (embedding_dim, embedding_dim))  # Value权重
        
        self.cached_input: Optional[np.ndarray] = None  # 缓存输入，用于反向传播
        
        # 为每个权重矩阵创建优化器
        self.optimizer_w_q = Adam((embedding_dim, embedding_dim))
        self.optimizer_w_k = Adam((embedding_dim, embedding_dim))
        self.optimizer_w_v = Adam((embedding_dim, embedding_dim))
    
    def compute_qkv(self, input_data: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        计算查询(Q)、键(K)、值(V)矩阵
        
        比喻理解：
        - Q (Query): "我想找什么信息？" - 当前词的查询向量
        - K (Key): "我能提供什么信息？" - 每个词的键向量  
        - V (Value): "我的具体信息内容" - 每个词的值向量
        
        就像在图书馆：
        - Q是你要查找的主题
        - K是每本书的标签/目录
        - V是书的实际内容
        """
        q = input_data @ self.w_q  # Q = 输入 × Q权重矩阵
        k = input_data @ self.w_k  # K = 输入 × K权重矩阵
        v = input_data @ self.w_v  # V = 输入 × V权重矩阵
        return q, k, v
    
    def attention(self, q: np.ndarray, k: np.ndarray, v: np.ndarray) -> np.ndarray:
        """
        计算注意力机制的核心逻辑
        
        步骤：
        1. 计算注意力分数：Q和K的相似度
        2. 应用因果掩码：防止看到"未来"的词
        3. 转换为注意力权重：用softmax归一化
        4. 加权求和：根据权重组合V的信息
        """
        # 缩放因子：防止分数过大导致梯度消失
        dk = np.sqrt(self.embedding_dim)
        
        # 计算注意力分数矩阵
        # scores[i][j] 表示第i个词对第j个词的关注程度
        k_t = k.T  # 转置K矩阵
        scores = (q @ k_t) / dk  # 计算相似度并缩放
        
        # 应用因果掩码（Causal Masking）
        # 防止模型在预测时"偷看"未来的词汇
        # 就像考试时不能看后面的题目来答前面的题
        seq_len = scores.shape[0]
        for i in range(seq_len):
            for j in range(i + 1, seq_len):
                scores[i, j] = -np.inf  # 将未来位置设为负无穷
        
        # 计算注意力权重
        weights = self.softmax(scores)  # 转换为概率分布
        
        # 加权求和：根据注意力权重组合信息
        return weights @ v
    
    def softmax(self, scores: np.ndarray) -> np.ndarray:
        """
        Softmax函数 - 把分数转换为概率分布
        
        输入：注意力分数矩阵
        输出：注意力权重矩阵（每行和为1）
        
        就像把考试分数转换为排名百分比
        """
        result = scores.copy()
        
        # 对每一行应用softmax
        for i in range(result.shape[0]):
            row = result[i, :]
            
            # 数值稳定性处理：减去最大值防止指数溢出
            max_val = np.max(row)
            exp_values = np.exp(row - max_val)
            sum_exp = np.sum(exp_values)
            
            # 归一化为概率分布
            result[i, :] = exp_values / sum_exp
        
        return result
    
    @staticmethod
    def softmax_backward(softmax_output: np.ndarray, grad_output: np.ndarray) -> np.ndarray:
        """
        Softmax的反向传播
        计算softmax函数的梯度，用于参数更新
        
        这是数学推导的结果，用于训练时的梯度计算
        """
        grad_input = softmax_output.copy()
        
        for i in range(grad_input.shape[0]):
            softmax_row = softmax_output[i, :]
            grad_out_row = grad_output[i, :]
            
            # 计算softmax梯度的数学公式
            dot = np.sum(softmax_row * grad_out_row)
            
            for j in range(grad_input.shape[1]):
                y_i = softmax_row[j]
                dy_i = grad_out_row[j]
                grad_input[i, j] = y_i * (dy_i - dot)
        
        return grad_input
    
    def layer_type(self) -> str:
        """返回层类型名称"""
        return "SelfAttention"
    
    def forward(self, input_data: np.ndarray) -> np.ndarray:
        """
        前向传播 - 处理输入数据
        
        步骤：
        1. 缓存输入（用于反向传播）
        2. 计算Q、K、V矩阵
        3. 执行注意力计算
        4. 添加残差连接
        
        残差连接：output = attention_result + input
        就像在原有信息基础上添加新的理解，而不是完全替换
        """
        self.cached_input = input_data.copy()  # 保存输入用于反向传播
        
        # 计算注意力
        q, k, v = self.compute_qkv(input_data)
        attention_output = self.attention(q, k, v)
        
        # 残差连接：保留原始信息，防止信息丢失
        return attention_output + input_data
    
    def backward(self, grads: np.ndarray, lr: float) -> np.ndarray:
        """
        反向传播 - 根据错误信号更新参数
        
        这个过程比较复杂，涉及链式法则和矩阵求导
        主要目的是计算每个参数应该如何调整来减少错误
        """
        # 获取缓存的输入和中间计算结果
        input_data = self.cached_input
        q = input_data @ self.w_q
        k = input_data @ self.w_k
        v = input_data @ self.w_v
        
        # 重新计算注意力分数和权重
        dk = self.w_q.shape[1]
        scale = np.sqrt(dk)
        scores = (q @ k.T) / scale
        
        # 应用因果掩码
        seq_len = scores.shape[0]
        for i in range(seq_len):
            for j in range(i + 1, seq_len):
                scores[i, j] = -np.inf
        
        attn_weights = self.softmax(scores)
        
        # 计算各部分的梯度
        grad_attn_weights = grads @ v.T
        grad_v = attn_weights.T @ grads
        
        # Softmax反向传播
        grad_scores = self.softmax_backward(attn_weights, grad_attn_weights)
        
        # 计算Q和K的梯度
        grad_q = grad_scores @ k
        grad_k = grad_scores.T @ q
        
        # 计算权重矩阵的梯度
        grad_w_q = input_data.T @ grad_q
        grad_w_k = input_data.T @ grad_k
        grad_w_v = input_data.T @ grad_v
        
        # 计算输入的梯度
        grad_input_attention = (
            grad_q @ self.w_q.T +
            grad_k @ self.w_k.T +
            grad_v @ self.w_v.T
        )
        
        # 添加残差连接的梯度
        grad_input = grad_input_attention + grads
        
        # 更新权重参数
        self.optimizer_w_q.step(self.w_q, grad_w_q, lr)
        self.optimizer_w_k.step(self.w_k, grad_w_k, lr)
        self.optimizer_w_v.step(self.w_v, grad_w_v, lr)
        
        return grad_input

# 使用示例和理解：
# 
# 假设输入句子："猫 坐在 垫子 上"
# 
# 1. 每个词都会生成Q、K、V向量
# 2. "坐在"的Q向量会与所有词的K向量计算相似度
# 3. 发现与"猫"和"垫子"的相似度较高
# 4. 根据相似度权重，融合"猫"和"垫子"的V向量信息
# 5. 让"坐在"这个词获得更丰富的上下文理解